"""
根据参考的 union_find ，完成以下功能的实现
    在代码中实现这些步骤：
    初始 a = [1,2,3,4,5],并将其添加到并查集里
    分别合并[1,2][3,5] [3,1]
    然后判断 2 5 是否为同一个集合
"""


class Node:
    def __init__(self, data, parent):
        self.data = data
        self.parent = parent


class UnionFind:
    def __init__(self, elements):
        self.elements = elements
        self.nodes = []
        for element in elements:
            node = Node(data=element, parent=None)
            self.nodes.append(node)

    def node(self, element):
        return self.nodes[self.elements.index(element)]

    def union(self, element1, element2):
        root1 = self.node(element1)
        root2 = self.node(element2)
        while root1.parent != None:
            root1 = root1.parent
        while root2.parent != None:
            root2 = root2.parent
        if root1.data < root2.data:
            root2.parent = root1
        else:
            root1.parent = root2

    def find(self, element):
        root = self.node(element)
        while root.parent != None:
            root = root.parent
        return root

    def is_same_set(self, element1, element2):
        return self.find(element1) == self.find(element2)


a = [1, 2, 3, 4, 5]
Set = UnionFind(a)
Set.union(1, 2)
Set.union(3, 5)
Set.union(3, 1)
print(Set.is_same_set(2, 5))
